import time
import argparse
import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

def initialize():
    parser = argparse.ArgumentParser(
        description="convert format between xvg,csv"
    )
    parser.add_argument(
        '-f',
        help="Name of inputfile with extension",
    )
    parser.add_argument(
        '-o',
        
        default="OutputFile",
        help="Name of output file with extension",
    )
    
    args = parser.parse_args()
    return args


def read_xvg(xaxis,yaxis,time_list,rmsd_list):
    with open(input_file, 'r') as fo:
        lines = fo.readlines()
    
    # xaxis,yaxis = "",""
    # time_list, rmsd_list = [],[]

    for line in lines:
        if line[0] == "#":
            continue
        if line[0] == "@":
            if "xaxis" in line:
                xaxis = line.strip().strip('"').split('"')[-1]
            if "yaxis" in line:
                yaxis = line.strip().strip('"').split('"')[-1]
    	
        if not line.startswith("@") and not line.startswith("#"):
            items = line.strip().split()
            time_list.append(float(items[0]))
            rmsd_list.append(float(items[1]))
    	   
    if len(time_list) != len(rmsd_list):
        print(len(time_list), len(rmsd_list),'\n', "Wrong in length time_list and rmsd_list")

def read_csv(x,y,y1):  # 从列表读数据
	df = pd.read_csv(input_file, sep=',',header=0)
	x = df['Time (ps)']
	y = df['RMSD (nm)']
	d = pd.Series(y)

	y1 = d.rolling(50).mean() # 计算滑动平均值


def save_csv(xaxis,yaxis,time_list,rmsd_list):
	  	
	with open(out_file, 'w') as fo :
		fo.write(xaxis + "," +yaxis+'\n')
		for i in range(len(time_list)):
			fo.write("{},{}\n".format(time_list[i], rmsd_list[i])) 	   


def plot_save(x,y,y1):
	plt.style.use('seaborn-darkgrid')  # 绘图风格fast，ggplot，classic,默认default
	plt.figure(figsize=(16,9))
	plt.plot(x, y,c='black')
	
	plt.plot(x,y1, c='red')
	font={"family":"SimHei","style":"normal","weight":"bold","color":"red","size":30}
	plt.title("RMSD",fontdict=font)
	plt.xlabel('Time(ps)',fontdict={'size':24})
	plt.ylabel('Rmsd(nm)',fontdict={'size':24})
	
	# plt.legend()
	plt.savefig(out_file,dpi=600,bbox_inches = 'tight')  # 一定要在show()之前执行



if __name__ == "__main__":
    start_time = time.time()
    args = initialize()
    if args.f is None:
        print("No input file given, exiting now.")
        exit(1)
    else:
        if os.path.exists(args.f):
            input_file = args.f
        else:
            print("Input file does not exist! (or cannot be found)")
            exit(0)
    if os.path.exists(args.o):
        print("Output file already exists. Exiting to prevent overwriting data.")
        exit(0)
    out_file = args.o

    xaxis,yaxis = "",""
    time_list,rmsd_list = [],[]

    x,y,y1 = [],[],[]

    if input_file.split('.')[1] == "xvg" and out_file.split('.')[1] == "csv": 
    	# input xvg，output csv
    	
    	read_xvg(xaxis,yaxis,time_list,rmsd_list)
    	save_csv(xaxis,yaxis,time_list,rmsd_list)
    	
    if input_file.split('.')[1] == "xvg" and out_file.split('.')[1] == "png":
    	# input xvg, output png
    	read_xvg(xaxis,yaxis,time_list,rmsd_list)
    	save_csv(xaxis,yaxis,time_list,rmsd_list)

    	plot_save(x,y,y1)

    if input_file.split('.')[1] == "csv" and out_file.split('.')[1] == "png":
    	read_csv(x,y,y1)
    	plot_save(x,y,y1)
    


    elapsed_time = time.time() - start_time
    print(time.strftime("%H:%M:%S", time.gmtime(elapsed_time)))    